// pkg/service/snapshotter.go
package service

import (
	"context"
	"fmt"
	"path/filepath"
	"sync"

	"github.com/containerd/containerd/v2/core/mount"
	"github.com/containerd/containerd/v2/core/snapshots"
	"github.com/onyx/snapshotter/pkg/cache"
	"github.com/onyx/snapshotter/pkg/config"
	"github.com/onyx/snapshotter/pkg/dedup"
	"github.com/onyx/snapshotter/pkg/fanotify"
	"github.com/onyx/snapshotter/pkg/metadata"
	"github.com/onyx/snapshotter/pkg/ublk"
)

// Snapshotter 实现 containerd 的 Snapshotter 接口
type Snapshotter struct {
	config      *config.Config
	metaDB      metadata.Store
	dedupEngine *dedup.Engine
	cacheLayer  *cache.Layer
	fanotify    *fanotify.Listener
	ublkBackend *ublk.Backend
	mu          sync.RWMutex
	snapshots   map[string]*SnapshotInfo
}

// SnapshotInfo 跟踪快照元数据
type SnapshotInfo struct {
	Key      string
	Parent   string
	Kind     snapshots.Kind
	Readonly bool
	Mounts   []mount.Mount
}

// NewSnapshotter 创建新的 Snapshotter 实例
func NewSnapshotter(cfg *config.Config) (*Snapshotter, error) {
	// 初始化元数据数据库
	metaDB, err := metadata.NewStore(cfg.DBBackend, filepath.Join(cfg.Root, "metadata"))
	if err != nil {
		return nil, fmt.Errorf("failed to init metadata store: %w", err)
	}

	// 初始化去重引擎
	dedupEngine, err := dedup.NewEngine(
		filepath.Join(cfg.Root, "chunks"),
		cfg.ChunkSize,
		metaDB,
	)
	if err != nil {
		return nil, fmt.Errorf("failed to init dedup engine: %w", err)
	}

	// 初始化缓存层
	cacheLayer := cache.NewLayer(filepath.Join(cfg.Root, "cache"), 10*1024*1024*1024) // 10GB

	// 初始化 fanotify 监听器
	var fanotifyListener *fanotify.Listener
	if cfg.FanotifyEnabled {
		fl, err := fanotify.NewListener()
		if err != nil {
			return nil, fmt.Errorf("failed to init fanotify: %w", err)
		}
		fanotifyListener = fl
	}

	// 初始化 ublk 后端
	var ublkBackend *ublk.Backend
	if cfg.UblkEnabled {
		ub, err := ublk.NewBackend(filepath.Join(cfg.Root, "ublk"))
		if err != nil {
			return nil, fmt.Errorf("failed to init ublk backend: %w", err)
		}
		ublkBackend = ub
	}

	return &Snapshotter{
		config:      cfg,
		metaDB:      metaDB,
		dedupEngine: dedupEngine,
		cacheLayer:  cacheLayer,
		fanotify:    fanotifyListener,
		ublkBackend: ublkBackend,
		snapshots:   make(map[string]*SnapshotInfo),
	}, nil
}

// Prepare 准备新的快照（通常用于创建临时层）
func (s *Snapshotter) Prepare(ctx context.Context, key string, parent string, opts ...snapshots.Opt) ([]mount.Mount, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if _, exists := s.snapshots[key]; exists {
		return nil, fmt.Errorf("snapshot %s already exists", key)
	}

	info := &SnapshotInfo{
		Key:      key,
		Parent:   parent,
		Kind:     snapshots.KindActive,
		Readonly: false,
	}

	// 根据 parent 构建 mounts
	// 如果有 parent，则使用 overlayfs 或 loop+overlay
	mounts := s.buildMounts(parent, key, false)
	info.Mounts = mounts

	s.snapshots[key] = info
	return mounts, nil
}

// View 创建只读视图
func (s *Snapshotter) View(ctx context.Context, key string, parent string, opts ...snapshots.Opt) ([]mount.Mount, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if _, exists := s.snapshots[key]; exists {
		return nil, fmt.Errorf("snapshot %s already exists", key)
	}

	info := &SnapshotInfo{
		Key:      key,
		Parent:   parent,
		Kind:     snapshots.KindView,
		Readonly: true,
	}

	mounts := s.buildMounts(parent, key, true)
	info.Mounts = mounts

	s.snapshots[key] = info
	return mounts, nil
}

// Commit 提交快照为不可变层
func (s *Snapshotter) Commit(ctx context.Context, name string, key string, opts ...snapshots.Opt) error {
	s.mu.Lock()
	_, exists := s.snapshots[key]
	s.mu.Unlock()

	if !exists {
		return fmt.Errorf("snapshot %s not found", key)
	}

	// 更新元数据
	s.mu.Lock()
	info := s.snapshots[key]
	info.Kind = snapshots.KindCommitted
	info.Readonly = true
	
	s.snapshots[name] = info
	delete(s.snapshots, key)
	s.mu.Unlock()

	// 触发去重：扫描文件系统，计算分块哈希
	if err := s.dedupEngine.Deduplicate(ctx, info.Mounts); err != nil {
		return fmt.Errorf("deduplication failed: %w", err)
	}

	return nil
}

// Remove 删除快照
func (s *Snapshotter) Remove(ctx context.Context, key string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	_, exists := s.snapshots[key]
	if !exists {
		return fmt.Errorf("snapshot %s not found", key)
	}

	// 清理相关的块存储、元数据
	// 注意：实际实现中需要根据key查找相关信息进行清理
	// 这里简化处理
	if err := s.dedupEngine.CleanupSnapshot(ctx, key); err != nil {
		return fmt.Errorf("failed to cleanup snapshot: %w", err)
	}

	delete(s.snapshots, key)
	return nil
}

// Stat 获取快照状态
func (s *Snapshotter) Stat(ctx context.Context, key string) (snapshots.Info, error) {
	s.mu.RLock()
	info, exists := s.snapshots[key]
	s.mu.RUnlock()

	if !exists {
		return snapshots.Info{}, fmt.Errorf("snapshot %s not found", key)
	}

	snapshotInfo := snapshots.Info{
		Name:   key,
		Parent: info.Parent,
		Kind:   info.Kind,
	}
	
	// 注意：Readonly 不是 snapshots.Info 的字段，需要通过 Kind 判断
	return snapshotInfo, nil
}

// Update 更新快照（例如改变 parent）
func (s *Snapshotter) Update(ctx context.Context, info snapshots.Info, opts ...snapshots.Opt) (snapshots.Info, error) {
	// 实现更新逻辑
	return info, nil
}

// Walk 遍历快照树
func (s *Snapshotter) Walk(ctx context.Context, fn snapshots.WalkFunc, filters ...string) error {
	s.mu.RLock()
	defer s.mu.RUnlock()

	for _, info := range s.snapshots {
		snapshotInfo := snapshots.Info{
			Name:   info.Key,
			Parent: info.Parent,
			Kind:   info.Kind,
		}
		
		if err := fn(ctx, snapshotInfo); err != nil {
			return err
		}
	}
	return nil
}

// Close 关闭 Snapshotter
func (s *Snapshotter) Close() error {
	if s.fanotify != nil {
		s.fanotify.Close()
	}
	if s.ublkBackend != nil {
		s.ublkBackend.Close()
	}
	if s.cacheLayer != nil {
		s.cacheLayer.Close()
	}
	if s.metaDB != nil {
		s.metaDB.Close()
	}
	return nil
}

// buildMounts 构建挂载点
// 如果有 parent，使用 overlayfs；否则返回块设备挂载
func (s *Snapshotter) buildMounts(parent, key string, readonly bool) []mount.Mount {
	mounts := make([]mount.Mount, 0)

	if parent != "" {
		// overlayfs 模式
		lowerdir := filepath.Join(s.config.Root, "snapshots", parent, "fs")
		upperdir := filepath.Join(s.config.Root, "snapshots", key, "fs")
		workdir := filepath.Join(s.config.Root, "snapshots", key, "work")

		mounts = append(mounts, mount.Mount{
			Type:   "overlay",
			Source: "overlay",
			Options: []string{
				fmt.Sprintf("lowerdir=%s", lowerdir),
				fmt.Sprintf("upperdir=%s", upperdir),
				fmt.Sprintf("workdir=%s", workdir),
			},
		})
	} else if s.ublkBackend != nil {
		// ublk 块设备模式
		devicePath := fmt.Sprintf("/dev/ublk%d", hashKey(key)%256)
		options := []string{}
		if readonly {
			options = append(options, "ro")
		}
		mounts = append(mounts, mount.Mount{
			Type:    "ext4",
			Source:  devicePath,
			Options: options,
		})
	}

	return mounts
}

// hashKey 简单哈希函数
func hashKey(s string) uint32 {
	h := uint32(0)
	for _, c := range s {
		h = h*31 + uint32(c)
	}
	return h
}